#!/bin/bash

# 指向昇腾软件包安装地址，导出环境变量
if [ ! $ASCEND_HOME_DIR ]; then
    if [ -d "$HOME/Ascend/ascend-toolkit/latest" ]; then
        export ASCEND_HOME_DIR=$HOME/Ascend/ascend-toolkit/latest
    else
        export ASCEND_HOME_DIR=/usr/local/Ascend/ascend-toolkit/latest
    fi
fi
source $ASCEND_HOME_DIR/bin/setenv.bash

# 获取当前的目录
CURRENT_DIR=$(
    cd $(dirname ${BASH_SOURCE:-$0})
    pwd
); cd $CURRENT_DIR

rm -rf src/build/* output/*.bin input/*.bin out/*

function get_cxx_env() {
    res=$(python3 -c "import torch" &> /dev/null || echo "torch_not_exist")
    if [ "$res" == "torch_not_exist" ]; then
        cxx_abi=1
    else
        if [ $(python3 -c 'import torch; print(torch.compiled_with_cxx11_abi())') == "True" ]; then
            cxx_abi=1
        else
            cxx_abi=0
        fi
    fi
    export CXX_ABI=$cxx_abi
}

declare -A VersionMap
VersionMap["Ascend310P1"]="Ascend310P1"
VersionMap["Ascend310P3"]="Ascend310P1"
VersionMap["Ascend910B1"]="Ascend910B1"
VersionMap["Ascend910B2"]="Ascend910B1"
VersionMap["Ascend910B3"]="Ascend910B1"
VersionMap["Ascend910B4"]="Ascend910B1"

declare -A SOCMap
SOCMap["Ascend310P1"]="ascend310p"
SOCMap["Ascend910B1"]="ascend910B1"

mkdir ./input
mkdir ./output
source $ASCEND_HOME_DIR/bin/setenv.bash
# 生成计算输入数据和对比用的真值数据
python golden_calc.py

function get_soc_name() {
    res=$(python3 -c 'import acl; print(acl.get_soc_name()[:11])')
    export SOC_VERSION=${VersionMap[$res]}
    soc=${VersionMap[$res]}
}
CORE_TYPE=VectorCore

# 编译ACLNN算子插件
function compile_plugin_aclnn_op() {
    cd ../../examples/atb_plugin/plugin_aclnn_operations
    mkdir -p ./build
    cd ./build
    cmake .. -DCMAKE_CXX_COMPILER=g++ -DCMAKE_SKIP_RPATH=TRUE -DCMAKE_CXX_STANDARD=17 -DCMAKE_CXX_STANDARD_REQUIRED=ON
    make

    cp libplugin_aclnn_op.so ../../out/libplugin_aclnn_op.so
    cd $CURRENT_DIR

    if [ $? != 0 ];then
        echo "make for plugin_aclnn_operations Successfully"
        exit 0
    fi
}

# 编译自定义算子插件
function compile_plugin_custom_op() {
    cd ../../examples/atb_plugin/plugin_custom_operations
    rm -rf build
    mkdir -p build; cd build;                                 
    cmake .. -DCMAKE_CXX_COMPILER=g++ -DCMAKE_SKIP_RPATH=TRUE -DCMAKE_CXX_STANDARD=17 -DCMAKE_CXX_STANDARD_REQUIRED=ON
    make

    cp libplugin_custom_op.so ../../out/libplugin_custom_op.so
    cd $CURRENT_DIR

    if [ $? != 0 ];then
        echo "make for plugin_custom_operations Successfully"
        exit 0
    fi
}

#编译自定义算子kernel
function compile_op_kernel() {
    cd ../../examples/atb_plugin/custom_op_kernel
    rm -rf build
    mkdir -p build; cd build;                                 
    get_soc_name
    cmake ..                                                 \
        -DASCEND_PRODUCT_TYPE=${SOCMap[$soc]}    \
        -DASCEND_CORE_TYPE=$CORE_TYPE                        \
        -DASCEND_RUN_MODE=ONBOARD                          \
        -DASCEND_INSTALL_PATH=$ASCEND_HOME_DIR
    cmake --build . --target ascendc_kernels
    cd -

    if [ $? -ne 0 ]; then
        echo "ERROR: compile op on failed!"
        return 1
    fi
    echo "INFO: compile op on npu succeed!"
    cp ./libs/libascendc_kernels.so ../out/libascendc_kernels.so
    cd $CURRENT_DIR
}

get_cxx_env
rm -rf ../../examples/atb_plugin/out/*
mkdir ../../examples/atb_plugin/out
compile_plugin_aclnn_op
compile_op_kernel
compile_plugin_custom_op

cd $CURRENT_DIR

cd ./src
mkdir ./build
cd ./build
cmake .. -DCMAKE_CXX_COMPILER=g++ -DCMAKE_SKIP_RPATH=TRUE -DCMAKE_CXX_STANDARD=17 -DCMAKE_CXX_STANDARD_REQUIRED=ON
cmake --build .

cd $CURRENT_DIR
export LD_LIBRARY_PATH=$CURRENT_DIR/../../examples/atb_plugin/out:$LD_LIBRARY_PATH
export ATB_WORKSPACE_MEM_ALLOC_GLOBAL=1

mkdir ./out
cd ./out
./main